class Generator(nn.Module):
def __init__(self, z_dim, img_channels, features_g, resize=False):
super(Generator, self).__init__()
self.resize = resize
if (self.resize == False):
self.gen = nn.Sequential(
self.layer(in_channels=z_dim, out_channels=features_g*64, kernel_size=4, stride=1, padding=0, bn=True), # 1x1
self.layer(in_channels=features_g*64, out_channels=features_g*32, kernel_size=4, stride=2, padding=1, bn=True), # 4x4
self.layer(in_channels=features_g*32, out_channels=features_g*16, kernel_size=4, stride=2, padding=1, bn=True), # 8x8
self.layer(in_channels=features_g*16, out_channels=features_g*8, kernel_size=4, stride=2, padding=1, bn=True), # 16x16
self.layer(in_channels=features_g*8, out_channels=features_g*4, kernel_size=4, stride=2, padding=1, bn=True), #32x32
self.layer(in_channels=features_g*4, out_channels=features_g*2, kernel_size=4, stride=2, padding=1, bn=True), # 64x64
nn.ConvTranspose2d(in_channels=features_g*2, out_channels=img_channels, kernel_size=4, stride=2, padding=1), # 256x256
nn.Tanh() # Nair & Hinton, 2010 -> mapira (normalizira) na vrijednosti pravih slika
)
else:
self.gen = nn.Sequential(
self.layer(in_channels=z_dim, out_channels=features_g*8, kernel_size=4, stride=1, padding=0, bn=True), # 1x1
self.layer(in_channels=features_g*8, out_channels=features_g*4, kernel_size=4, stride=2, padding=1, bn=True), # 4x4
self.layer(in_channels=features_g*4, out_channels=features_g*2, kernel_size=4, stride=2, padding=1, bn=True), # 8x8
self.layer(in_channels=features_g*2, out_channels=features_g, kernel_size=4, stride=2, padding=1, bn=True), # 16x16
# self.layer(in_channels=features_g, out_channels=features_g, kernel_size=4, stride=2, padding=1, bn=True), # 32x32
nn.ConvTranspose2d(in_channels=features_g, out_channels=img_channels, kernel_size=4, stride=2, padding=1, bias=False), # 64x64
nn.Tanh()
)
def layer(self, in_channels, out_channels, kernel_size, stride, padding, bn=True):
if bn:
return nn.Sequential(
nn.ConvTranspose2d(in_channels=in_channels, out_channels=out_channels, kernel_size=kernel_size, stride=stride, padding=padding, bias=False),
nn.BatchNorm2d(num_features=out_channels),
nn.ReLU(True)
)
else:
return nn.Sequential(
nn.ConvTranspose2d(in_channels=in_channels, out_channels=out_channels, kernel_size=kernel_size, stride=stride, padding=padding, bias=False),
nn.ReLU(True)
)
def forward(self, input):
return self.gen(input)
class Discriminator(nn.Module):
def __init__(self, img_channels, features_d, resize=False):
super(Discriminator, self).__init__()
self.resize = resize
if (self.resize == False):
self.disc = nn.Sequential(
# 3x256x256
self.layer(in_channels=img_channels, out_channels=features_d, kernel_size=4, stride=2, padding=1, bn=False), # fdx128x128
self.layer(in_channels=features_d, out_channels=features_d*2, kernel_size=4, stride=2, padding=1, bn=True), # 64x64
self.layer(in_channels=features_d*2, out_channels=features_d*4, kernel_size=4, stride=2, padding=1, bn=True), # 32x32
self.layer(in_channels=features_d*4, out_channels=features_d*8, kernel_size=4, stride=2, padding=1, bn=True), #16x16
self.layer(in_channels=features_d*8, out_channels=features_d*16, kernel_size=4, stride=2, padding=1, bn=True), #8x8
self.layer(in_channels=features_d*16, out_channels=features_d*32, kernel_size=4, stride=2, padding=1, bn=True), # 4x4
nn.Conv2d(in_channels=features_d*32, out_channels=1, kernel_size=4, stride=2, padding=0), # 1x1
# nn.Sigmoid()
)
else:
self.disc = nn.Sequential(
# 3x64x64
self.layer(in_channels=img_channels, out_channels=features_d, kernel_size=4, stride=2, padding=1, bn=False), # 64x64
self.layer(in_channels=features_d, out_channels=features_d*2, kernel_size=4, stride=2, padding=1, bn=True), # 32x32
self.layer(in_channels=features_d*2, out_channels=features_d*4, kernel_size=4, stride=2, padding=1, bn=True), # 16x16
self.layer(in_channels=features_d*4, out_channels=features_d*8, kernel_size=4, stride=2, padding=1, bn=True), #8x8
nn.Conv2d(in_channels=features_d*8, out_channels=1, kernel_size=4, stride=2, padding=0, bias=False), # 4x4
# nn.Sigmoid()
)
def layer(self, in_channels, out_channels, kernel_size, stride, padding, bn=True):
if bn:
return nn.Sequential(
nn.Conv2d(in_channels=in_channels, out_channels=out_channels, kernel_size=kernel_size, stride=stride, padding=padding, bias=False),
nn.BatchNorm2d(num_features=out_channels), # Batch normalizacija
nn.LeakyReLU(0.1, inplace=True)
)
return nn.Sequential(
nn.Conv2d(in_channels=in_channels, out_channels=out_channels, kernel_size=kernel_size, stride=stride, padding=padding, bias=False),
nn.LeakyReLU(0.1, inplace=True)
)
def forward(self, input):
return self.disc(input)
